#include <netdb.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <strings.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>

#include "psd_conn_udt.h"
#include "psd_protocols.h"
#include "psd_logger.h"
#include "psd_config.h"
#include "psd_tpool.h"
#include "psd_modules.h"
#include "psd_settings.h"

#include "compat.h"

struct udt_listener_args {
	UDTSOCKET socket;
	pthread_t accept_thread;
};

int psd_proto_udt_init();
static psdConn *psd_proto_udt_connect(const char *depot_id, psdSettings *settings);
static psdListener *psd_proto_udt_setup_listener(const char *listener_id, psdSettings *settings, int one_shot, listener_cb callback, void *arg);
static int psd_proto_udt_start_listener(psdListener *listener);
static int psd_proto_udt_stop_listener(psdListener *listener);
static void psd_proto_udt_free_listener(psdListener *listener);
static void *psd_udt_accept_handler(void *v);

static psdProtocolHandler psd_udt_handler = {
	.connect = psd_proto_udt_connect,
	.setup_listener = psd_proto_udt_setup_listener,
	.name = "udt"
};

#ifndef EMBED_UDT
static psdModule psd_udt_module = {
	.desc = "UDT Handler Module",
	.dependencies = "udt", 			// XXX: gross hack
	.init = psd_proto_udt_init
};

psdModule *module_info() {
	return &psd_udt_module;
}
#endif

int udt_def_use_ipv6 = 0;
int udt_def_incoming_port = 5006;
int udt_def_outgoing_port = 5006;
int udt_def_mss = 1500;
int udt_def_incoming_recv_bufsize = 0;
int udt_def_incoming_send_bufsize = 0;
int udt_def_outgoing_recv_bufsize = 0;
int udt_def_outgoing_send_bufsize = 0;

int psd_proto_udt_init() {

	if (psd_add_protocol_handler(&psd_udt_handler)) {
		psd_err(0, "couldn't add protocol handler");
		goto error_exit;
	}

	return 0;

error_exit:
	return -1;
}

static psdConn *psd_proto_udt_connect(const char *hostname, psdSettings *settings) {
	struct addrinfo hints;
	struct addrinfo *depot_addrs, *nexthop;
	char port_str[10];
	int port;
	int retval;
	psdConn *ret_conn;
	int connected;
	UDTSOCKET new_sd;
	int recv_bufsize;
	int send_bufsize;
	int mss;
	int temp;
	int use_ipv6;
	char *bind_interface;
	struct hostent *he;
	struct sockaddr_storage sa;
	char* udt_congestion_ctrl; 
	int udt_enable_custom_cc;

	// grab the udt settings from the given policy
	if (psd_settings_get_int_2(settings, "udt", "mss", &mss) != 0) {
		mss = udt_def_mss;
	}
	if (psd_settings_get_int_2(settings, "udt", "port", &port) != 0) {
		port = udt_def_outgoing_port;
	}

	if (psd_settings_get_int_2(settings, "udt", "recv_bufsize", &recv_bufsize) != 0) {
		recv_bufsize = udt_def_outgoing_recv_bufsize;
	}

	if (psd_settings_get_int_2(settings, "udt", "send_bufsize", &send_bufsize) != 0) {
		send_bufsize = udt_def_outgoing_send_bufsize;
	}

	if (psd_settings_get_bool_2(settings, "udt", "use_ipv6", &use_ipv6) != 0) {
		use_ipv6 = udt_def_use_ipv6;
	}

	if (psd_settings_get_2(settings, "udt", "interface", &bind_interface) != 0) {
		bind_interface = NULL;
	}

	if (psd_settings_get_int_2(settings, "udt", "custom_cc", &udt_enable_custom_cc) != 0) {
	  udt_enable_custom_cc=0;
        }

	if(udt_enable_custom_cc == 1)
	  if (psd_settings_get_2(settings, "udt", "ccc_name", &udt_congestion_ctrl) != 0) {
	    udt_congestion_ctrl = NULL;
	  }

	psd_info(1, "connecting to %s/%d", hostname, port);

	snprintf(port_str, 10, "%d", port);

	bzero(&hints, sizeof(struct addrinfo));

	if (use_ipv6) {
		hints.ai_family = AF_INET6;
	} else {
		hints.ai_family = AF_INET;
	}

	hints.ai_socktype = SOCK_STREAM;
	hints.ai_flags |= AI_PASSIVE;

	retval = getaddrinfo(hostname, port_str, &hints, &depot_addrs);
	if (retval != 0) {
		psd_err(1, "lookup of \"%s\" failed", hostname);
		goto error_exit;
	}

	if (bind_interface != NULL) {
		he = gethostbyname((const char *) bind_interface);
		if (he != NULL) {
			bzero((void *)&sa, sizeof(struct sockaddr_storage));
			((struct sockaddr *)&sa)->sa_family = he->h_addrtype;
			if (he->h_addrtype == AF_INET) {
				memcpy (&(((struct sockaddr_in *) &sa)->sin_addr), he->h_addr_list[0], he->h_length);
				((struct sockaddr_in *) &sa)->sin_port = 0;
			} else {
				memcpy (&(((struct sockaddr_in6 *) &sa)->sin6_addr), he->h_addr_list[0], he->h_length);
				((struct sockaddr_in6 *) &sa)->sin6_port = 0;
			}
		} else {
			bind_interface = NULL;
		}
	}

	connected = 0;

	for(nexthop = depot_addrs; nexthop != NULL && !connected; nexthop = nexthop->ai_next) {
		new_sd = udt_socket(nexthop->ai_family, nexthop->ai_socktype, nexthop->ai_protocol);
		if (new_sd < 0) {
			psd_warn(10, "socket() failed");
			continue;
		}

		if (bind_interface != NULL) {
			udt_bind(new_sd, (struct sockaddr *)&sa, sizeof(sa));
		}

		if (recv_bufsize > 0) {
			udt_setsockopt(new_sd, SOL_SOCKET, SO_RCVBUF, &(recv_bufsize), sizeof(int));
		}

		if (send_bufsize > 0) {
			udt_setsockopt(new_sd, SOL_SOCKET, SO_SNDBUF, &(send_bufsize), sizeof(int));
		}
		
		psd_info(5, "setting UDT_MSS to %d", mss);
		udt_setsockopt(new_sd, 0, UDT_MSS, &mss, sizeof(int));

		if(udt_enable_custom_cc==1) {
		  udt_setsockopt(new_sd,SOL_SOCKET, UDT_CUSTOM_CC_ENABLED, &udt_enable_custom_cc, sizeof(int));
		  psd_info(0,"UDT Custom Congestion Control Algorithm is enabled");
		  
		  if( udt_congestion_ctrl != NULL) {
		    udt_setsockopt(new_sd,SOL_SOCKET, UDT_CCN, udt_congestion_ctrl, strlen(udt_congestion_ctrl));
		    psd_info(0,"UDT Congestion Control Algorithm is %s",udt_congestion_ctrl);
		  }
		}

		if (udt_connect(new_sd, nexthop->ai_addr, nexthop->ai_addrlen) != 0) {
			psd_warn(5, "connect() failed");
			udt_close(new_sd);
			continue;
		}

		connected = 1;
	}

	freeaddrinfo(depot_addrs);

	if (!connected) {
		psd_err(0, "couldn't connect to \"%s\"", hostname);
		goto error_exit;
	}
	
	ret_conn = psd_conn_udt_alloc(new_sd);
	if (!ret_conn) {
		psd_err(10, "couldn't allocate socket structure");
		goto error_exit_sd;
	}

	psd_info(1, "connected to %s", hostname);

	return ret_conn;

error_exit_sd:
	udt_close(new_sd);
error_exit:
	return NULL;
}

static psdListener *psd_proto_udt_setup_listener(const char *listener_id, psdSettings *settings, int one_shot, listener_cb callback, void *arg) {
	psdListener *new_listener;
	struct udt_listener_args *udt_args;

	new_listener = psd_listener_alloc();
	if (!new_listener)
		goto error_exit;

	new_listener->id = strdup(listener_id);
	if (!new_listener->id)
		goto error_exit_listener;

	new_listener->settings = psd_settings_duplicate(settings);
	if (!new_listener->settings)
		goto error_exit_listener;

	udt_args = malloc(sizeof(struct udt_listener_args));
	if (!udt_args)
		goto error_exit_listener;

	new_listener->status = LISTENER_STOPPED;
	new_listener->callback = callback;
	new_listener->arg = arg;
	new_listener->protocol = "UDT";
	new_listener->proto_private = udt_args;
	new_listener->start = psd_proto_udt_start_listener;
	new_listener->stop = psd_proto_udt_stop_listener;
	new_listener->free = psd_proto_udt_free_listener;

	return new_listener;

error_exit_listener:
	psd_listener_free(new_listener);
error_exit:
	return NULL;
}

static int psd_proto_udt_start_listener(psdListener *listener) {
	struct addrinfo *ai;
	int psd_udt_family;
	int i;
	int use_ipv6;
	char *bind_interface;
	int recv_bufsize;
	int send_bufsize;
	int mss;
	int temp;
	int port;
	struct hostent *he;
	struct sockaddr_storage sa;
	UDTSOCKET socket;
	int family;
	struct udt_listener_args *udt_args = listener->proto_private;
	char* udt_congestion_ctrl;
        int udt_enable_custom_cc;


	if (listener->status != LISTENER_STOPPED)
		return -1;

	// grab the udt settings from the given policy
	if (psd_settings_get_int_2(listener->settings, "udt", "mss", &mss) != 0) {
		mss = udt_def_mss;
	}
	if (psd_settings_get_int_2(listener->settings, "udt", "port", &port) != 0) {
		port = udt_def_incoming_port;
	}
	if (psd_settings_get_int_2(listener->settings, "udt", "recv_bufsize", &recv_bufsize) != 0) {
		recv_bufsize = udt_def_incoming_recv_bufsize;
	}

	if (psd_settings_get_int_2(listener->settings, "udt", "send_bufsize", &send_bufsize) != 0) {
		send_bufsize = udt_def_incoming_send_bufsize;
	}

	if (psd_settings_get_2(listener->settings, "udt", "interface", &bind_interface) != 0) {
		bind_interface = NULL;
	}

	if (psd_settings_get_bool_2(listener->settings, "udt", "use_ipv6", &use_ipv6) != 0) {
		use_ipv6 = udt_def_use_ipv6;
	}

	if (psd_settings_get_int_2(listener->settings, "udt", "custom_cc", &udt_enable_custom_cc) != 0) {
	  udt_enable_custom_cc=0;
        }
	
        if(udt_enable_custom_cc == 1)
          if (psd_settings_get_2(listener->settings, "udt", "ccc_name", &udt_congestion_ctrl) != 0) {
            udt_congestion_ctrl = NULL;
          }

	if (bind_interface != NULL) {
		he = gethostbyname((const char *) bind_interface);
		if (he != NULL) {
			bzero((void *)&sa, sizeof(struct sockaddr_storage));
			((struct sockaddr *)&sa)->sa_family = he->h_addrtype;
			if (he->h_addrtype == AF_INET) {
				memcpy (&(((struct sockaddr_in *) &sa)->sin_addr), he->h_addr_list[0], he->h_length);
				((struct sockaddr_in *) &sa)->sin_port = 0;
			} else {
				memcpy (&(((struct sockaddr_in6 *) &sa)->sin6_addr), he->h_addr_list[0], he->h_length);
				((struct sockaddr_in6 *) &sa)->sin6_port = 0;
			}
		} else {
			bind_interface = NULL;
		}
	}

	if (use_ipv6) {
		family = AF_INET6;
	} else {
		family = AF_INET;
	}

	socket = udt_socket(family, SOCK_STREAM, 0);
	if (socket == INVALID_UDTSOCK) {
		psd_warn(10, "socket() failed");
		goto error_exit_socket;
	}

	psd_info(5, "setting listener UDT_MSS to %d", mss);
	udt_setsockopt(socket, 0, UDT_MSS, &mss, sizeof(int));


	if(udt_enable_custom_cc==1) {
	  udt_setsockopt(socket,SOL_SOCKET, UDT_CUSTOM_CC_ENABLED, &udt_enable_custom_cc, sizeof(int));
	  psd_info(0,"UDT Custom Congestion Control Algorithm is enabled");

	  if( udt_congestion_ctrl != NULL) {
	    udt_setsockopt(socket,SOL_SOCKET, UDT_CCN, udt_congestion_ctrl, strlen(udt_congestion_ctrl));
	    psd_info(0,"UDT Congestion Control Algorithm is %s",udt_congestion_ctrl);
	  }
	}


	if (bind_interface != NULL) {
		udt_bind(socket, (struct sockaddr *)&sa, sizeof(sa));
	} else {
		struct sockaddr_in srv_addr;

		srv_addr.sin_family = AF_INET;
		srv_addr.sin_port = htons(port);
		srv_addr.sin_addr.s_addr = INADDR_ANY;
		memset(&(srv_addr.sin_zero), '\0', 8);

		if (udt_bind(socket, (struct sockaddr*)&srv_addr, sizeof(srv_addr)) != 0) {
			psd_err(0, "bind failed: %s", udt_getlasterror());
			goto error_exit_socket;
		}
	}

	if (recv_bufsize > 0) {
		udt_setsockopt(socket, SOL_SOCKET, SO_RCVBUF, &(recv_bufsize), sizeof(int));
	}

	if (send_bufsize > 0) {
		udt_setsockopt(socket, SOL_SOCKET, SO_SNDBUF, &(send_bufsize), sizeof(int));
	}

	if (udt_listen(socket, 32) != 0) {
		psd_err(0, "couldn't set socket to listening state: %s", udt_getlasterror());
		goto error_exit_socket;
	}

	udt_args->socket = socket;

	if (pthread_create(&udt_args->accept_thread, NULL, psd_udt_accept_handler, (void *) listener) != 0) {
		psd_err(0, "couldn't execute socket handler thread: %s", strerror(errno));
		goto error_exit_socket;
	}

	listener->status = LISTENER_RUNNING;
	
	psd_info(0, "Started UDT listener: %s", listener->id);

	return 0;

error_exit_socket:
	udt_close(socket);
error_exit:
	return -1;
}

static int psd_proto_udt_stop_listener(psdListener *listener) {
	struct udt_listener_args *udt_args = listener->proto_private;
	int i, n;

	if (listener->status != LISTENER_RUNNING) {
		psd_err(0, "tried to stop a stopped listener: %d", listener->status);
		return -1;
	}

	if ((n = pthread_cancel(udt_args->accept_thread)) < 0){
		psd_err(0, "pthread_cancel failed: %d, %s", n, strerror(errno));
		return -1;
	}

	if ((n = pthread_join(udt_args->accept_thread, NULL)) < 0){
		psd_err(0, "pthread_join failed: %d, %s", n, strerror(errno));
		return -1;
	}

	udt_close(udt_args->socket);

	listener->status = LISTENER_STOPPED;

	return 0;
}

static void psd_proto_udt_free_listener(psdListener *listener) {
	struct udt_listener_args *udt_args = listener->proto_private;

	if (listener->status != LISTENER_STOPPED)
		psd_proto_udt_stop_listener(listener);

	// once it is stopped, all the resources associated with the
	// listener(thread, sockets, etc) should be free, so we can just free
	// the struct
	free(udt_args);	
}

static void *psd_udt_accept_handler(void *v) {
	psdConn *new_conn;
	fd_set socket_set;
	int n, i, high;
	UDTSOCKET sd;
	psdListener *listener;
	struct udt_listener_args *udt_args;
	int old_state;

	listener = v;
	udt_args = listener->proto_private;

	pthread_detach(pthread_self());

	while(1) {
		struct sockaddr_storage sa;
		SOCKLEN_T sa_len = sizeof(struct sockaddr_storage);

		sd = udt_accept(udt_args->socket, (struct sockaddr *) &sa, &sa_len);
		if (sd == INVALID_UDTSOCK) {
			psd_err(1, "accept failed: %s", udt_getlasterror());
			continue;
		}

		new_conn = psd_conn_udt_alloc(sd);
		if (!new_conn) {
			psd_err(1, "psd_conn_alloc_socket() failed: %s", strerror(errno));
			udt_close(sd);
			continue;
		}

		new_conn->id = strdup(new_conn->description);

		psd_info(0, "new incoming connection: %s", new_conn->description);

		new_conn->settings = psd_settings_duplicate(listener->settings);

		n = listener->callback(listener, new_conn, listener->arg);
		if (n == 0 && listener->one_shot)
			goto out;

		pthread_setcancelstate(PTHREAD_CANCEL_ENABLE, &old_state);
	}

out:
	udt_close(udt_args->socket);

	listener->status = LISTENER_STOPPED;

	return NULL;
}
